{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e07826a-4e11-445c-a23d-67793e9852c3",
   "metadata": {},
   "source": [
    "### 0. 执行前说明\n",
    "- 需要修改的参数\n",
    "    - root：sar_aircraft项目原始数据的根目录，和current_path(也就是代码所在路径)是同级的\n",
    "    - yolo_datasets_path ：yolo对应数据集的路径\n",
    "- 原始数据目录结构\n",
    "    - 目录结构：\n",
    "            - current_path\n",
    "            - sar_aircraft\n",
    "                - JPEGImages\n",
    "                - Annotations\n",
    "      注：原始数据需要解压"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73757ffa-3119-45aa-9fda-b03b13f66beb",
   "metadata": {},
   "source": [
    "### 1. 读取分类文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0badd845-4f04-4fb2-b42a-c7130b3e262a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import shutil # 处理文件和目录\n",
    "from xml.etree import ElementTree\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "02528fec-c90f-41a9-b766-3470a1295cd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(16463, 2)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv(filepath_or_buffer=\"output.csv\")\n",
    "data.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a7dab09-57ed-4002-a01c-a5d856a71e13",
   "metadata": {},
   "source": [
    "### 2. 定义字典"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8203e1a1-a111-493f-b40a-29592bb9e848",
   "metadata": {},
   "outputs": [],
   "source": [
    "label2idx = {\"A330\": 0, \n",
    "             \"A320/321\": 1, \n",
    "             \"A220\": 2, \n",
    "             \"ARJ21\":3, \n",
    "             \"Boeing737\": 4, \n",
    "             \"Boeing787\":5, \n",
    "             \"other\":6}\n",
    "# 反向标签\n",
    "idx2label = {idx: label for label, idx in label2idx.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c49dd79b-263a-481a-83b2-b6266200bcc9",
   "metadata": {},
   "source": [
    "### 3. 划分A330类的数据集"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34f54aba-e98a-44d2-a3e8-58fd22da6165",
   "metadata": {},
   "source": [
    "思路：\n",
    "- 找出所有包含A330的图像作为正样本， 并随机找出相同数量的不包含A330的图像做负样本\n",
    "- 分为train（75%），val（20%）， test（5%）\n",
    "- test做最终验证，做为考题"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37405f8b-d63a-491c-98f4-9c0ab56022f0",
   "metadata": {},
   "source": [
    "#### 3.1 建立A330二分类的目录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "938ebac3-dd8d-49a5-beec-3c774823fa41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 原始数据的目录\n",
    "current_path = os.getcwd()\n",
    "root = os.path.join(current_path, \"sar_aircraft_source\")\n",
    "root_images_path = os.path.join(root, \"JPEGImages\")\n",
    "root_labels_path = os.path.join(root, \"Annotations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "02be5272-4d49-4cdf-8f3f-b62070a4a79a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "项目目录已重新创建\n"
     ]
    }
   ],
   "source": [
    "# 数据划分：定义训练集和验证集占原始数据的比例\n",
    "train_ratio = 0.75\n",
    "val_ratio = 0.2\n",
    "test_ratio = 0.05\n",
    "\n",
    "# 1. 定义YOLO的路径和数据集的路径\n",
    "yolo_datasets_path = os.path.join(\"D:\\\\\", \"datasets\")\n",
    "sar_aircraft_path = os.path.join(yolo_datasets_path, \"sar_aircraft\")\n",
    "classification_path = os.path.join(sar_aircraft_path, \"two_classification\")\n",
    "\n",
    "# 2. 在yolo的dataset中创建数据集目录，有的话直接删除，重新创建\n",
    "try:\n",
    "    shutil.rmtree(classification_path)\n",
    "    os.makedirs(classification_path)\n",
    "    print(\"项目目录已重新创建\")\n",
    "except FileNotFoundError:\n",
    "    os.makedirs(classification_path)\n",
    "    print(\"项目目录已创建\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c69fa1fd-9ad4-4303-b09b-922f8cf376b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# yolo需要的目录结构\n",
    "images_labels_list = [\"images\", \"labels\"]\n",
    "train_val_test_list = [\"train\", \"val\", \"test\"]\n",
    "\n",
    "# 创建各个分类的目录\n",
    "for i in range(7):\n",
    "    cls_path = os.path.join(classification_path, str(i))\n",
    "    os.makedirs(cls_path)\n",
    "    for path_0 in images_labels_list:\n",
    "        images_labels_path = os.path.join(cls_path, path_0)\n",
    "        os.makedirs(images_labels_path)\n",
    "        for path_1 in train_val_test_list:\n",
    "            train_val_test_path = os.path.join(images_labels_path, path_1)\n",
    "            os.makedirs(train_val_test_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7785a734-0c2a-44d0-ab3b-dfb54c495ae8",
   "metadata": {},
   "source": [
    "#### 3.2 找出A330的图像的正反例"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "450ce4a1-4568-42b4-a31e-60ccc5130b6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>文件名</th>\n",
       "      <th>第一列内容</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0000001.txt</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0000001.txt</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0000002.txt</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0000002.txt</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0000002.txt</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16458</th>\n",
       "      <td>0004368.txt</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16459</th>\n",
       "      <td>0004368.txt</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16460</th>\n",
       "      <td>0004368.txt</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16461</th>\n",
       "      <td>0004368.txt</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16462</th>\n",
       "      <td>0004368.txt</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>16463 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "               文件名  第一列内容\n",
       "0      0000001.txt      5\n",
       "1      0000001.txt      4\n",
       "2      0000002.txt      5\n",
       "3      0000002.txt      2\n",
       "4      0000002.txt      2\n",
       "...            ...    ...\n",
       "16458  0004368.txt      4\n",
       "16459  0004368.txt      2\n",
       "16460  0004368.txt      2\n",
       "16461  0004368.txt      2\n",
       "16462  0004368.txt      2\n",
       "\n",
       "[16463 rows x 2 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "01fd2234-60fd-4fe1-b679-f33a850044f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "类别0的图像 正样本 数量有：290\n",
      "类别不包含0的图像数量有：4078\n",
      "类别不包含0的图像 负样本 数量有：290\n"
     ]
    }
   ],
   "source": [
    "# 找出包含类别0的图像\n",
    "pos_img_0_file = list(set(data[data[\"第一列内容\"] == 0][\"文件名\"]))\n",
    "# 包含类别0图像的数量\n",
    "pos_img_0_num = len(pos_img_0_file)\n",
    "print(f\"类别0的图像 正样本 数量有：{pos_img_0_num}\")\n",
    "\n",
    "# 找出不含包类别0的所有图像图像\n",
    "neg_img_0_all_fille =set()\n",
    "for i in data[\"文件名\"].values:\n",
    "    if i not in pos_img_0_file:\n",
    "        neg_img_0_all_fille.add(i)\n",
    "print(f\"类别不包含0的图像数量有：{len(neg_img_0_all_fille)}\")\n",
    "\n",
    "# 从neg_img_0_all随机找出pos_img_0_num个负样本\n",
    "neg_img_0_file = random.sample(list(neg_img_0_all_fille), pos_img_0_num)\n",
    "print(f\"类别不包含0的图像 负样本 数量有：{len(neg_img_0_file)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "078a9ba9-3016-4994-b213-80418ccd72fc",
   "metadata": {},
   "source": [
    "#### 3.3 切分数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bb08ea90-f861-4c9c-aff0-38119f1da210",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 转list\n",
    "pos_img_0_file = list(pos_img_0_file)\n",
    "neg_img_0_file = list(neg_img_0_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7d176c13-ac4d-4e7f-9ebf-c554a9e2740b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_train_files, pos_test_files, neg_train_files, neg_test_files = train_test_split(pos_img_0_file, neg_img_0_file, test_size=test_ratio, shuffle=True, random_state=42)\n",
    "pos_train_files, pos_val_files, neg_train_files, neg_val_files = train_test_split(pos_train_files, neg_train_files, test_size=(val_ratio / (val_ratio + train_ratio)), shuffle=True, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0163d066-6489-44d4-994a-7858988ebee4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "217 58 15\n",
      "217 58 15\n"
     ]
    }
   ],
   "source": [
    "# # 查看形状\n",
    "print(len(pos_train_files), len(pos_val_files), len(pos_test_files))\n",
    "print(len(neg_train_files), len(neg_val_files), len(neg_test_files))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47537019-5e65-4661-a168-3bbe949a6de5",
   "metadata": {},
   "source": [
    "#### 3.4 转换标签文件格式，从xml转换到txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e40c5f3d-e54c-4acf-a746-b17ca591c50e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def xml_transfer_yolo(source_file_name, object_folder, file):\n",
    "    \"\"\"\n",
    "        定义转换格式的函数\n",
    "        把xml（voc）格式转化为txt（yolo）格式\n",
    "    \"\"\"\n",
    "    tree = ElementTree.parse(source=source_file_name)\n",
    "    root = tree.getroot()\n",
    "\n",
    "    # 图像的高度和宽度\n",
    "    img_width = int(root.find(path=\"size\").find(path=\"width\").text)\n",
    "    img_height = int(root.find(path=\"size\").find(path=\"height\").text)\n",
    "\n",
    "    # 转化标注信息\n",
    "    # finaly_file_name = os.path.splitext(file)[0] + \".txt\"\n",
    "    with open(file=os.path.join(object_folder, file), mode=\"w\", encoding=\"utf8\") as f:\n",
    "        # 遍历每个目标\n",
    "        # print(finaly_file_name)\n",
    "        for obj in root.findall(path=\"object\"):\n",
    "            name = obj.find(path=\"name\").text\n",
    "            cls_id = label2idx.get(name) if name == \"A330\" else \"1\"\n",
    "            xmin = int(obj.find(path=\"bndbox\").find(path=\"xmin\").text)\n",
    "            ymin = int(obj.find(path=\"bndbox\").find(path=\"ymin\").text)\n",
    "            xmax = int(obj.find(path=\"bndbox\").find(path=\"xmax\").text)\n",
    "            ymax = int(obj.find(path=\"bndbox\").find(path=\"ymax\").text)\n",
    "\n",
    "            \"\"\"\n",
    "                开始标注转化\n",
    "                - yolo需要的是相对坐标，也就是相对于原始图像的比例，这样当图像resize到不同\n",
    "                尺寸时，相对坐标所标注的框位置不变，可以更好的应对不同尺寸的图像\n",
    "            \"\"\"\n",
    "            # 1. 中心点坐标\n",
    "            # print(type(xmin), type(xmax), type(img_width))\n",
    "            x_center = round(number=(xmin + xmax) / 2 / img_width, ndigits=6)\n",
    "            y_center = round(number=(ymin + ymax) / 2 / img_height, ndigits=6)\n",
    "            box_width = round(number=(xmax - xmin) / img_width, ndigits=6)\n",
    "            box_height = round(number=(ymax - ymin) / img_height, ndigits=6)\n",
    "            print(cls_id, x_center, y_center, box_width, box_height, sep=\" \", end=\"\\n\", file=f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fae049e-2f96-47f1-b3af-817669fd45b8",
   "metadata": {},
   "source": [
    "### 3.5 复制img和label到yolo数据集目录下"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6960fb4b-fa43-4464-b601-a364645b10f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def copy_files(files, images_path, labels_path):\n",
    "    \"\"\"\n",
    "        复制文件到yolo的数据集目录下\n",
    "            - files: 需要复制的图像文件[train_files, val_files, test_files]\n",
    "            - images_path: 需要复制到的图像路径 \n",
    "            - labels_path: 需要复制到的标签路径\n",
    "    \"\"\"\n",
    "    for file in files:\n",
    "        # 复制图像\n",
    "        shutil.copy(os.path.join(root_images_path, os.path.splitext(file)[0] + '.jpg'), images_path)\n",
    "        # 复制标签\n",
    "        label_file = os.path.splitext(file)[0] + '.xml'\n",
    "        xml_transfer_yolo(source_file_name=os.path.join(root_labels_path, label_file), object_folder=labels_path, file=file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "af3ef7d5-5993-4a7f-bae8-40b747e5d668",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 训练集\n",
    "copy_files(files=pos_train_files, images_path=os.path.join(classification_path, \"0\", \"images\", \"train\"), labels_path=os.path.join(classification_path, \"0\", \"labels\", \"train\"))\n",
    "copy_files(files=neg_train_files, images_path=os.path.join(classification_path, \"0\", \"images\", \"train\"), labels_path=os.path.join(classification_path, \"0\", \"labels\", \"train\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1551e97d-43e2-4bcb-9300-df82b173c6b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 验证集\n",
    "copy_files(files=pos_val_files, images_path=os.path.join(classification_path, \"0\", \"images\", \"val\"), labels_path=os.path.join(classification_path, \"0\", \"labels\", \"val\"))\n",
    "copy_files(files=neg_val_files, images_path=os.path.join(classification_path, \"0\", \"images\", \"val\"), labels_path=os.path.join(classification_path, \"0\", \"labels\", \"val\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "18cb1675-7c69-4e8e-a82c-c21b262da069",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 训练集\n",
    "copy_files(files=pos_test_files, images_path=os.path.join(classification_path, \"0\", \"images\", \"test\"), labels_path=os.path.join(classification_path, \"0\", \"labels\", \"test\"))\n",
    "copy_files(files=neg_test_files, images_path=os.path.join(classification_path, \"0\", \"images\", \"test\"), labels_path=os.path.join(classification_path, \"0\", \"labels\", \"test\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68dc8a39-3530-4d5a-b0a5-61ee02f9340d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cf3deeb-7c83-4c0c-b4f3-035c3e18c1c0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
